import os
import re
import json
import jieba
import argparse

def remove_pun(text):
    text = text.lower().replace(',', ' ').replace('.', ' ').replace('?', ' ').replace('!', ' ').replace('，', '')\
        .replace('。','').replace('？', '').replace(':', ' ').replace(';', ' ').replace('<', ' ').replace('>', ' ')\
        .replace('/', ' ').replace('！', '').replace('"', '').replace('[',' ').replace(']',' ').replace('|','').replace('~','')\
        .replace('noise','').replace('(','').replace(')','').replace('ENG','').replace('eng','')
    return text

def get_args():
    parser = argparse.ArgumentParser(description="""deal crowd file""")
    parser.add_argument('-i','--input', type=str, required=True, default='',
                        help='input txt path')
    parser.add_argument('-o', '--outdir', type=str, required=True, default='',
                        help='output dir')
    args = parser.parse_args()
    return args


def get_path(path,pattern='.crowd'):
    result_dict = {}
    for root, dirs, files in os.walk(path):
        for f in files:
            #fpath = os.path.abspath(os.path.join(root, f))
            #print(fpath)
            if f.endswith(pattern):
                result_dict[f] = os.path.abspath(os.path.join(root, f))
            #print(result_dict)
            #path = os.path.split(fpath)
        return result_dict

def deal_crowd(input_file,text_file,seg_file,word_file):
    wav_urls = {}
    crowd_dict = get_path(input_file, pattern='.crowd')
    for fs in crowd_dict:
        #print(fs)
        with open(crowd_dict[fs], 'r', encoding='utf-8') as fin:
            #wav_info = {'req_file': input_file}
            wav_info = {'req_file': fs}
            for item in fin:
                j = json.loads(item)
                line = item.strip()
                line_info = eval(line)
                p_id, wave_name, wave_suf, data = [j[x] for x in
                                               ['p_id', 'wave_name', 'wave_suf', 'data']]
                wav_info[wave_name] = [p_id, wave_name, wave_suf, data]
                # seg word
                text = data[0]['mark_info']
                mark_text = str(text[0]['mark_text'])
                sen = remove_pun(mark_text)
                print(sen)

                if sen != '':
                   word = ' '.join(jieba.cut(sen, HMM=False))
                   text = re.sub(' +', ' ', word)
                   word_file.write(text + '\n')

                   wav_urls[wave_name] = wave_name + '.' + wave_suf + str(data[0]['start_time']) + str(data[0]['end_time']) + str(data[0]['text'])

                   text = wave_name + '.' + wave_suf + '\t ' + text
                   print(text)

                   segments = wave_name  + '.' + wave_suf + '\t ' + str(data[0]['start_time']) + '\t ' + str(data[0]['end_time'])

                   text_file.write(text + '\n')
                   seg_file.write(segments + '\n')
    seg_file.close()
    text_file.close()
    word_file.close()

    
def main():
    args = get_args()
    input_file = os.path.join(args.input)
    outdir = args.outdir
    seg_fout = open(os.path.join(outdir, 'segments'), 'w')
    text_fout = open(os.path.join( outdir, 'text'), 'w')
    word_fout = open(os.path.join( outdir, 'word'), 'w')
    deal_crowd(input_file,text_fout,seg_fout,word_fout)


if __name__ == "__main__":
    main()
